import os
import sqlite3
from llama_index.core import Document
from llama_index.core import VectorStoreIndex
from pprint import pprint
from llama_index.readers.database import DatabaseReader

reader = DatabaseReader(
    # 获取数据库类型，默认为mysql
    scheme=os.getenv("DB_SCHEME", "mysql"),
    # 获取数据库主机地址，默认为localhost
    host=os.getenv("DB_HOST", "localhost"),
    # 获取数据库端口，默认为3306
    port=os.getenv("DB_PORT", "3306"),
    # 获取数据库用户名，默认为root
    user=os.getenv("DB_USER", "root"),
    # 获取数据库密码，默认为123456
    password=os.getenv("DB_PASS", "root"),
    # 获取数据库名称，默认为jobs
    dbname=os.getenv("DB_NAME", "jobs"),
)

documents = reader.load_data(query="SELECT * FROM users ")
for i, doc in enumerate(documents):
    print(f"{i+1}:{doc.text}")
